library(librarian)
shelf(tidybulk, GEOquery, ggpubr, tidyverse, tidySummarizedExperiment)

quick_process_bulk <- function(tdb_in, group, skip_scale = F, keep_abundant = T,
                               min.count = 10) {
  tdb_out <- tdb_in |>
    aggregate_duplicates() |>
    identify_abundant(factor_of_interest = {{ group }}, minimum_counts = min.count)
  
  if (skip_scale) {
    tdb_out <- tdb_out |> scale_abundance()
  }
  
  if (keep_abundant) {
    tdb_out <- tdb_out |> keep_abundant()
  }
  
  tdb_out
}

plot_qc_bulk <- function(tdb_in, scaled_abundance, group) {
  g1 <- tdb_in |>
    ggplot(aes({{ scaled_abundance }}, color= {{ group }})) +
    geom_density() +
    scale_x_log10()

  g2 <- tdb_in |>
    reduce_dimensions(method="PCA", .dims = 2) |>
    pivot_sample() |>
    ggplot(aes(PC1, PC2, color={{ group }})) +
    geom_point()
  
  print(g1 + g2)
  
  tdb_in |>
    reduce_dimensions(method="PCA", .dims = 2) |>
    pivot_sample()
}

download_ncbi_counts <- function(acc) {
  message('INFO: This function only works for human data on GEO!')
  ncbi_anno <-
  "https://www.ncbi.nlm.nih.gov/geo/download/?format=file&type=rnaseq_counts&type=rnaseq_counts&file=Human.GRCh38.p13.annot.tsv.gz" |>
    read_tsv(col_select = 1:2)
  
  target_file <-
  "https://www.ncbi.nlm.nih.gov/geo/download/?format=file&type=rnaseq_counts&acc=GSE124676&file=GSE124676_raw_counts_GRCh38.p13_NCBI.tsv.gz" |>
    str_replace_all('GSE\\d+', acc)
  
  message('INFO: Downloading ', target_file, ' ...')
  
  target_file |>
    read_tsv() |>
    left_join(ncbi_anno) |>
    relocate(Symbol) |>
    select(-GeneID)
}

download_ncbi_meta <- function(acc, simple = TRUE) {
  ncbi_meta <- getGEO(acc) |>
    pluck(1)
  
  ncbi_meta <- ncbi_meta@phenoData@data |>
    as_tibble() |>
    select(1:2, ends_with(':ch1'))
  
  if(simple){return(ncbi_meta[,1:2])} else {return(ncbi_meta)}
}

calc_tpm <- function(tdb_in, sample, abundance) {
  tdb_in |>
    summarize(total.count = sum({{ abundance }}), .by = {{ sample }}) |>
    right_join(tdb_in) |>
    mutate(tpm = {{ abundance }} / total.count * 1e6)
}

geo_limma <- function(eset, gene_col = 'ILMN_Gene', use_vooma = TRUE,
                      complete_case = TRUE, force_normalize = FALSE) {
  require(limma)
  fvarLabels(eset) <- fvarLabels(eset) |>
    make.names()
  ex <- exprs(eset)
  
  expr_quantiles <- ex |>
    quantile(c(0., 0.25, 0.99, 1.0), na.rm = T) |>
    as.numeric()
  
  ## if 99% quantile > 100 or (range > 50 and 25% > 0)
  LogC <- (expr_quantiles[3] > 100) ||
    (expr_quantiles[4] - expr_quantiles[1] > 50 && expr_quantiles[2] > 0)
  
  ## replace negative value with NaN before log2
  if (LogC) {
    message('Automatic log tranformation...')
    ex[which(ex <= 0)] <- NaN
    exprs(eset) <- log2(ex)
  }
  
  if (force_normalize) {
    message('Performing normalizeBetweenArrays...')
    exprs(eset) <- normalizeBetweenArrays(exprs(eset))
  }
  
  if (complete_case) {
    message('Removing cases with NA exprs...')
    eset <- eset[complete.cases(exprs(eset)), ]
  }
  
  design <- model.matrix(~group + 0, eset)
  
  if (use_vooma) {
    message('Apply vooma...')
    v <- vooma(eset, design)
    
    v$genes <- fData(eset) # attach gene annotations
    
    message('Fit linear model...')
    fit <- lmFit(v)
  } else {
    message('Vooma not applied. Fit linear model...')
    fit <- lmFit(eset, design)
  }

  # set up contrasts of interest and recalculate model coefficients
  cts <- str_c(colnames(design), collapse = '-')
  message(str_glue('Compute DEGs for {cts}...'))
  cont.matrix <- makeContrasts(contrasts=cts, levels=design)
  
  res <- fit |>
    contrasts.fit(cont.matrix) |>
    eBayes(0.01) |>
    topTable(adjust="fdr", number = Inf) |>
    as_tibble() |>
    select(ID, adj.P.Val, logFC) |>
    arrange(adj.P.Val)
  
  result <- try({
    res <- gene.annot <- eset@featureData@data |>
      select(all_of(c('ID', gene_col))) |>
      right_join(res) |>
      dplyr::rename('gene' = any_of(gene_col)) |>
      as_tibble()
  }, silent = TRUE)
  
  if (inherits(result, "try-error")) {
    warning("Gene annotation not found, return only probe ID.")
    res
  } else {
    res
  }
}

pluck_geo <- function(acc = NULL, file = NULL, annot_gpl = FALSE) {
  eset <- getGEO(acc, filename = file, AnnotGPL = annot_gpl)

  if(is.null(file)){
    if(length(eset) == 1) {
      eset <- eset[[1]]
    } else {
      warning('Find more than one file for this GEO!')
      return(eset)
    }
  }
  
  fvarLabels(eset) <- fvarLabels(eset) |> make.names()
  varLabels(eset) <- varLabels(eset) |> make.names()
  eset
}
